from typing import Optional
import torch
import numpy as np
import cv2


class CamBase:
    """
    基本仿射相机模型，不含外参数矩阵，能够实现从三维坐标点输出一张图像的转换。
    """

    def __init__(
        self,
        f_x: float,
        f_y: Optional[float],
        u_0: float,
        v_0: float,
        dx: float,
        dy: Optional[float],
        h: float,
        w: float,
        **kwargs,
    ) -> None:
        """
        Arguments:
         - `f_x` : x轴有效焦距，单位是毫米
         - `f_y` : y轴有效焦距，单位是毫米
         - `u_0` : x轴方向上主点与原点的偏移，单位是像素
         - `v_0` : y轴方向上主点与原点的偏移，单位是像素
         - `dx` : x轴上单位像素对应的毫米数
         - `dy` : y轴上单位像素对应的毫米数
         - `h` : 相机在x轴方向上的分辨率
         - `w` : 相机在y轴方向上的分辨率
        """
        self._K = torch.tensor(
            [[1 / dx, 0, u_0], [0, 1 / dy, v_0], [0, 0, 1]],
            dtype=torch.float32,
            device="cuda",
        )
        self._F = torch.tensor(
            [[f_x, 0, 0], [0, f_y, 0], [0, 0, 1]], dtype=torch.float32, device="cuda"
        )
        self._K = self._K @ self._F
        self._h = h
        self._w = w

    def __str__(self):
        """
        用于打印输出的字符串
        """
        return (
            f"-------Here is something about camera:\n"
            f"The inner parameter matrix is:\n{self.K}\n"
            f"The height of camera is {self.h}\n"
            f"The width of camera is {self.w}\n"
            f"The END of camera-------"
        )

    def output(self):
        return self.K

    @property
    def K(self):
        """
        内参数矩阵
        """
        return self._K

    @property
    def h(self):
        return self._h

    @property
    def w(self):
        return self._w

    def _photo(self, p_s: torch.Tensor, l_s: torch.Tensor) -> np.ndarray:
        """
        对相机坐标系投影，将某个给定的三维立体信息的点投影到成像平面。
        Arguments:
         - `p_s`: 一系列三维点，按列放置在一起成为一个矩阵
         - `l_s`: 一系列三维点的像素值，按列放置在一起成为一个矩阵
        """
        # 求得向量化的图像坐标，按行拉平
        image = torch.zeros(self.h, self.w, dtype=torch.uint8, device="cuda").flatten()
        # 将像素值赋给图像
        # 这一步忽略了空间中不同位置投影到同一个像素点处的强度，默认使用最后一个点强度代替了其他所有
        p_s_proj = (p_s[:, 0:2].T / torch.vstack([p_s[:, 2]] * 2)).to(torch.int64)
        ind = (
            (p_s_proj[0, :] >= 0)
            & (p_s_proj[0, :] < self.w)
            & (p_s_proj[1, :] >= 0)
            & (p_s_proj[1, :] < self.h)
        )
        flat_ind = p_s_proj[0, ind] * self.w + p_s_proj[1, ind]
        image[flat_ind] = l_s[ind]
        # 这一步是时间花得最多的，可以使用GPU加速
        # for i, cond in tqdm.tqdm(enumerate(ind.cpu().tolist())):
        #     if cond:
        #         image[p_s_proj[0, i], p_s_proj[1, i]] = l_s[i]
        img_cv = image.numpy(force=True)
        img_filtered = cv2.medianBlur(img_cv.reshape([self.h, self.w]), 3)
        return (
            (img_filtered - np.min(img_filtered))
            / (np.max(img_filtered) - np.min(img_filtered))
            * 255
        ).astype(np.uint8)
